package com.example.demo001_request_limit;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author sherry
 * @description
 * @date Create in 2020/4/10
 * @modified By:
 */

@RestController
@Slf4j
public class RequestController {

    /**
     * 单机网关限流使用一个ConcurrentHashMap存储bucket
     * 如果是分布式集群限流，需要使用Redis等分布式解决方案
     */
    private static final Map<String, Bucket> LOCAL_CACHE = new ConcurrentHashMap<>();

    /**
     * 桶容量，即token最大数量
     */
    private int capacity = 100000;

    /**
     * 每次Token补充量
     */
    private int refillTokens = 100000;

    /**
     * 补充Token的时间间隔
     */
    private Duration refillduration = Duration.ofSeconds(30);

    private Bucket createNewBucket() {
        Refill refill = Refill.intervally(refillTokens, refillduration);
        Bandwidth limit = Bandwidth.classic(capacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }

    @GetMapping("/test1")
    public ResponseEntity test1() {

        String key = "key";
        Bucket bucket = LOCAL_CACHE.computeIfAbsent(key, k -> createNewBucket());
        log.info("key:{}，令牌桶可用Token数量{}", key, bucket.getAvailableTokens());
        if (bucket.tryConsume(1)) {
            System.out.println("执行业务逻辑");
        } else {
            System.out.println("不执行业务逻辑");
        }

        return ResponseEntity.ok("结束");
    }

}
